You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
make CycleIterator save and restore its epoch through Fabric checkpoints
create the pretraining iterator before checkpoint loading and reuse it in fit
keep old checkpoints compatible by loading the existing state first, then restoring the optional iterator state from Fabric's remainder
cover mid-epoch and exact epoch-boundary resume with a real LitData StreamingDataLoader
Why
The current litdata==0.2.59 pin already restores the dataloader cursor correctly: after a mid-epoch checkpoint it yields the remaining tail once and then starts a complete new epoch. However, pretrain still constructs CycleIterator after loading and resets its epoch counter to zero.
This addresses the remaining epoch-tracking problem described in #1712 without changing dataloader cursor semantics. The boundary handling is needed because a restored StreamingDataLoader can begin the next epoch itself before CycleIterator observes StopIteration.
Validation
unmodified baseline: the real Fabric checkpoint regression fails because CycleIterator contains an unpicklable generator, and the pretraining resume path has no iterator state
CI note: the tests.yaml / Lit Job (n/a) check failed, but GitHub only exposes a generic “Job Failed” result with no annotations, and the Lightning job log requires separate authentication. The new regression tests in this PR are CPU tests and are not collected by this lane because it runs with RUN_ONLY_CUDA_TESTS=1. The same lane has also failed at a similar runtime on unrelated recent diffs (for example #2309 and #2312), so I do not want to change unrelated code without a traceback.
Could a maintainer please rerun the n/a lane once and approve the fork CPU workflow so the new regression tests run in the standard CI? Locally, pytest tests/test_pretrain.py tests/test_utils.py -q passed (59 passed, 9 skipped), and all applicable pre-commit hooks passed. I will inspect any concrete failing node/traceback before making further code changes.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changed
CycleIteratorsave and restore its epoch through Fabric checkpointsfitStreamingDataLoaderWhy
The current
litdata==0.2.59pin already restores the dataloader cursor correctly: after a mid-epoch checkpoint it yields the remaining tail once and then starts a complete new epoch. However,pretrainstill constructsCycleIteratorafter loading and resets its epoch counter to zero.This addresses the remaining epoch-tracking problem described in #1712 without changing dataloader cursor semantics. The boundary handling is needed because a restored
StreamingDataLoadercan begin the next epoch itself beforeCycleIteratorobservesStopIteration.Validation
CycleIteratorcontains an unpicklable generator, and the pretraining resume path has no iterator statepytest tests/test_pretrain.py tests/test_utils.py -q: 59 passed, 9 skippedgit diff --checkpassed